from PySide6.QtCore import QObject, Signal, Slot, QThread, Qt
from PySide6.QtGui import QIcon
from PySide6.QtWidgets import QTreeWidgetItem, QTreeWidget
from collections import defaultdict
from dbService.db_service import get_model
import logging


class TreeBuilder(QObject):
    finished = Signal(list)  # 线程完成后发射纯数据结构

    def __init__(self, data_list, root_id=0):
        super().__init__()
        self.data_list = data_list
        self.root_id = root_id
        self.running = True

    def stop(self):
        self.running = False

    def run(self):
        try:
            children_map = defaultdict(list)
            for fid, fparentid, fnumber, fname in self.data_list:
                children_map[fparentid].append((fid, fnumber, fname))

            def build_data(parent_id):
                if not self.running:
                    return []
                result = []
                for fid, fnumber, fname in children_map.get(parent_id, []):
                    result.append({
                        "fid": fid,
                        "fnumber": fnumber,
                        "fname": fname,
                        "children": build_data(fid)
                    })
                return result

            tree_data = build_data(self.root_id)
            self.finished.emit(tree_data)
        except Exception as e:
            logging.exception("TreeBuilder 运行异常: %s", e)


class AsyncTreeLoader(QObject):
    def __init__(self, sql: str, treeWidget: QTreeWidget, root_id=0,
                 folder_icon: QIcon = None, leaf_icon: QIcon = None):
        super().__init__()
        self.sql = sql
        self.treeWidget = treeWidget
        self.root_id = root_id
        self.folder_icon = folder_icon
        self.leaf_icon = leaf_icon

        self.thread = QThread()
        self.treeBuilder = None

    def start(self):
        self.treeWidget.clear()
        self.treeWidget.addTopLevelItem(QTreeWidgetItem(['正在加载...', '']))
        self.treeWidget.setHeaderLabels(self.treeWidget.HEADER_LABELS)

        self.treeWidget.itemExpanded.connect(self.adjust_columns)

        model = get_model(self.sql)

        data_list = [
            (
                model.index(row, 0).data(),
                model.index(row, 1).data(),
                model.index(row, 2).data(),
                model.index(row, 3).data()
            )
            for row in range(model.rowCount())
        ]

        self.treeBuilder = TreeBuilder(data_list, self.root_id)
        self.treeBuilder.moveToThread(self.thread)

        self.thread.started.connect(self.treeBuilder.run)
        self.treeBuilder.finished.connect(self.on_finished)
        self.treeBuilder.finished.connect(self.thread.quit)
        self.treeBuilder.finished.connect(self.treeBuilder.deleteLater)
        self.thread.finished.connect(self.thread.deleteLater)

        self.thread.start()

    @Slot(list)
    def on_finished(self, tree_data):
        self.treeWidget.clear()

        def add_items(parent_item, children):
            for node in children:
                item = QTreeWidgetItem([node["fnumber"], node["fname"]])
                item.setData(0, Qt.UserRole, node["fid"])
                item.setFlags(item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
                item.setCheckState(0, Qt.CheckState.Unchecked)

                # 设置图标（在主线程执行）
                if node["children"]:
                    if self.folder_icon:
                        item.setIcon(0, self.folder_icon)
                else:
                    if self.leaf_icon:
                        item.setIcon(0, self.leaf_icon)

                parent_item.addChild(item)
                add_items(item, node["children"])

        add_items(self.treeWidget.invisibleRootItem(), tree_data)

    @Slot(QTreeWidgetItem)
    def adjust_columns(self, item):
        for i in range(self.treeWidget.columnCount()):
            self.treeWidget.resizeColumnToContents(i)

    def stop(self):
        if self.thread.isRunning():
            self.treeBuilder.stop()
            self.thread.quit()
            self.thread.wait()
